import json
from large_model.mode_interface.model_interface import ChatInterface
from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain_core.prompts import MessagesPlaceholder, ChatPromptTemplate
from utils.file_dir import get_doc_md_dir
from utils.yaml_utils import read_template_yaml_all
from large_model.prompts_template import system_prompt, system_message, case_system_prompt, interface_document_prompt, \
    function_document_prompt
from large_model.model_tool.model_utils import save_test_cases_to_file, save_xmind_cases_to_file, \
    read_and_extract_interfaces, read_document_content
from large_model.prompt_template import system_messages
from large_model.model_tool.case_model import InitialTestCaseGenerator, TestCaseEnhancer, FunctionalTestCaseGenerator, \
    FunctionalTestCaseReviewer
from utils.data_utils import read_interface_datas
from langchain_community.document_loaders import UnstructuredFileLoader


class Executors:

    def __init__(self):
        # 初始化工具
        self.tools = [save_test_cases_to_file,
                      save_xmind_cases_to_file,
                      read_and_extract_interfaces,
                      read_document_content,
                      InitialTestCaseGenerator(),
                      TestCaseEnhancer(),
                      FunctionalTestCaseGenerator(),
                      FunctionalTestCaseReviewer(),  # 添加评审工具
                      ]
        # 防止影响工具输出，使用chat_model_tool
        self.llm = ChatInterface().chat_model_tool()

    # 接口测试用例， 单接口测试用例编写
    def interface_agent_main(self, executor_user_request, temp=False):
        if temp:
            systems_messages = system_messages + system_prompt + interface_document_prompt
        else:
            systems_messages = system_messages + system_prompt

        # 接口文档模板
        prompt = ChatPromptTemplate.from_messages([
            ("system", systems_messages),
            ("human", "用户请求：{input}"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ])

        # 7. 创建代理
        agent = create_tool_calling_agent(self.llm, self.tools, prompt)
        agent_executor = AgentExecutor(
            agent=agent,
            tools=self.tools,
            # 开启智能体日志
            verbose=True,
            max_iterations=20,
            max_execution_time=1200,
            handle_parsing_errors=True
        )

        #  5. 转义模板中的所有花括号
        executor_user_request = str(executor_user_request).replace('{', '{{').replace('}', '}}')
        # 9. 执行代理
        result = agent_executor.invoke({"input": executor_user_request})
        print(result["output"])

    # 功能测试用例
    def function_agent_main(self, executor_user_request, temp=False):
        # 1.根据输入的公共功能模块，生成测试用例
        # 2.根据读取的需求文档，在提取用户提到的功能模块，生成测试用例
        if temp:
            systems_messages = system_messages + system_message + function_document_prompt
        else:
            systems_messages = system_messages + system_message

        # 功能文档模板
        prompt = ChatPromptTemplate.from_messages([
            ("system", systems_messages),
            ("human", "用户请求：{input}"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ])

        # 7. 创建代理
        agent = create_tool_calling_agent(self.llm, self.tools, prompt)
        agent_executor = AgentExecutor(
            agent=agent,
            tools=self.tools,
            # 开启智能体日志
            verbose=True,
            max_iterations=20,
            max_execution_time=1200,
            handle_parsing_errors=True
        )

        #  5. 转义模板中的所有花括号
        executor_user_request = str(executor_user_request).replace('{', '{{').replace('}', '}}')
        # 9. 执行代理
        result = agent_executor.invoke({"input": executor_user_request})
        print(result["output"])

    # 没有文档的测试用例编写
    def agent_main(self, executor_user_request):
        systems_messages = system_messages + system_message
        # 6. 创建提示模板
        prompt = ChatPromptTemplate.from_messages([
            ("system", systems_messages),
            ("human", "用户请求：{input}"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ])

        # 7. 创建代理
        agent = create_tool_calling_agent(self.llm, self.tools, prompt)
        agent_executor = AgentExecutor(
            agent=agent,
            tools=self.tools,
            # 开启智能体日志
            verbose=True,
            max_iterations=40,
            max_execution_time=1200,
            handle_parsing_errors=True
        )

        # 9. 执行代理
        result = agent_executor.invoke({"input": executor_user_request})
        print(result["output"])

    def interface_mian(self):
        # 针对接口进行测试用例编写
        api_doc = read_interface_datas()
        # 串联的接口业务自己写， 或者辅助进行编写，也可以
        Executors().interface_agent_main(
            f"请你根据接口文档信息{api_doc},按照以下格式, 对GetSeatChannelConf 接口进行【单接口】的测试用例的编写并保存")

    def interface_document_mian(self):
        # 主要是针对接口文档进行测试用例的编写，直接在调用工具进行编写， 接口文档或者需求文档，都放到model_read_file下
        Executors().interface_agent_main(
            "要求：1.接口文档路径为:input_file=接口文档.docx ,请你对GetSeatChannelConf 接口进行【单接口】的测试用例的编写并保存",
            temp=True)

    def function_mian(self):
        # 功能柜测试用例编写
        loader = UnstructuredFileLoader(get_doc_md_dir("王者aipass问答.docx"))
        documents = loader.load()
        # 2. 直接使用文档内容
        document_content = documents[0].page_content
        # 串联的接口业务自己写， 或者辅助进行编写，也可以---默认开始测试用例评审
        Executors().function_agent_main(
            f"请你根据功能需求文档信息{document_content},进行功能用例的编写，"
            f"要求："
            f"1.不需要调用read_document_content工具，读取文档内容。请根据提供的信息进行编写测试用例"
            f"2.使用工具：FunctionalTestCaseReviewer，对编写生成的测试用例评审，只需要评审一次即可"
            f"3.使用工具: save_xmind_cases_to_file 保存为Xmind文件，其中input_text为Xmind文件的内容，output_file为Xmind文件的路径需要大模型根据功能模块生成生成")

    # 根据用户输入的公共功能模块，生成测试用例
    def case_main(self):
        temp = """
        请你根据以下要求编写测试用例：
        需要编写测试用例的功能或者需求:
        说明：特定助手号下，只要用户提问次数>X，后续的答案都带可点击的人工入口， 点击深度思考需要正常进行吧问题进行深度思考，如果没有深度思考给出的进人工入口，超过多少次就需要进人工可以点击进入人工
        1.正常情况下：如果没有开启R1深度思考入口，就推送【如未能解决问题，可尝试换个问法或点击人工】
        2.一旦用户提问次数>X（X支持by助手号来配置），就改为推送【如果觉得无法解决，可以点击进行深度思考，或点击人工】 或【如未能解决问题，可尝试换个问法或点击人工】（具体看有没有开启R1深度思考入口）
        """
        Executors().agent_main(temp + case_system_prompt)


if __name__ == '__main__':
    # 执行接口测试用例生成---这几个比较快
    # Executors().interface_mian()

    # 这个比较慢
    # Executors().interface_document_mian()

    # 执行功能测试用例生成--根据输入的公共功能模块，生成测试用例
    Executors().case_main()

    # 根据功能测试用例的描述，进行编写测试用例
    # Executors().function_mian()

    # 针对功能用例进行测试用例的编写
    # Executors().function_agent_main("要求：1.需求文档路径为:file_path=王者aipass问答.docx , 2.请你进行测试用例的编写，并保存为xmind文件", temp=True)
